package cn.wind.xboot.framework.filter;

import org.apache.commons.io.output.TeeOutputStream;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.PrintWriter;

/**
 * <p>Title: ResponseWrapper</p>
 * <p>Description: TODO</p>
 *
 * @author xukk
 * @version 1.0
 * @date 2018/8/9
 */
public class ResponseWrapper extends HttpServletResponseWrapper {
    private ByteArrayOutputStream _stream;
    private PrintWriter _pw;
    private ServletOutputStream _out;
    private TeePrintWriter _tee;
    private long id;


    public ResponseWrapper(HttpServletResponse response) throws IOException {
        super(response);
        init();
    }

    public ResponseWrapper(Long requestId, HttpServletResponse response) throws IOException {
        super(response);
        this.id = requestId;
        init();
    }

    private void init() throws IOException {
        _stream = new ByteArrayOutputStream();
        _pw = new PrintWriter(_stream);

    }

    /**
     * 覆盖getWriter()方法，将字符流缓冲到本地
     */
    @Override
    public PrintWriter getWriter() throws IOException {
        if(_tee==null){
           _tee= new TeePrintWriter(super.getWriter(), _pw);
            return _tee;
        }
        return _tee;
    }

    /**
     * 覆盖getOutputStream()方法，将字节流缓冲到本地
     */
    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        if(_out==null){
          _out =new ServletOutputStream() {
                @Override
                public boolean isReady() {
                    return false;
                }

                @Override
                public void setWriteListener(WriteListener writeListener) {

                }

                private TeeOutputStream tee = new TeeOutputStream(ResponseWrapper.super.getOutputStream(),_stream);

                @Override
                public void write(int b) throws IOException {
                    tee.write(b);
                }
            };
          return _out;
        }
        return _out;
    }

    public byte[] toByteArray() throws IOException {
        if(null!=_out){
            _out.flush();
            _out.close();
        }
        if(null!=_tee){
            _tee.flush();
            _tee.close();
        }
        if(null!=_stream){
            _stream.flush();
            _stream.close();
        }
        return _stream.toByteArray();
    }

    public long getId() {
        return id;
    }

}
